
#include "xdns-handler.h"
#include "xdns-address-query.h"
#include "xdns-data.h"
#include "xdns-decoder.h"
#include "xdns-header-info.h"
#include "xdns-host-entry-generator.h"
#include "xdns-init.h"
#include "xdns-name-query.h"
#include "xdns-protocol.h"
#include "xdns-question-info.h"
#include "xdns-resource-record-info.h"
#include "xdns-server-manager.h"
#include "xdns-server-state.h"
#include <arpa/inet.h>
#include <assert.h>
#include <string.h>
#include <limits>
#include <sstream>

using namespace std;

XdnsHandler::XdnsHandler()
{
}

XdnsHandler::~XdnsHandler()
{
}

XdnsError
XdnsHandler::Init(const XdnsConfig& config,
                  XdnsHandler** handler)
{
    if (NULL == handler)
    {
        return XE_Inval;
    }

    XdnsError error = XE_Ok;
    vector<XdnsServerState*> serverStates;
    XdnsServerManager* serverManager = NULL;
    do
    {
        error = XdnsHandlerImpl::InitServerStates(config.m_customServers,
                                                  serverStates);
        if (error != XE_Ok)
        {
            break;
        }
        assert(!(serverStates.empty()));

        error = XdnsServerManager::Init(serverStates, &serverManager);
        if (error != XE_Ok)
        {
            break;
        }

        // The server manager takes the ownership of the server states.
        serverStates.clear();

        XdnsHandlerImpl* handlerImpl = new (nothrow) XdnsHandlerImpl(serverManager);
        if (NULL == handlerImpl)
        {
            error = XE_NoMem;
            break;
        }

        // The XDNS handler takes the ownership of the server manager.
        serverManager = NULL;

        *handler = handlerImpl;
    }
    while (0);

    if (XE_Ok == error)
    {
        return XE_Ok;
    }

    // Destroy the temporary variables for the returned error.
    if (serverManager != NULL)
    {
        XdnsServerManager::Release(serverManager);
        serverManager = NULL;
    }

    vector<XdnsServerState*>::iterator iter = serverStates.begin();
    vector<XdnsServerState*>::iterator end = serverStates.end();
    for (; iter != end; ++iter)
    {
        XdnsServerState::Release(*iter);
    }
    serverStates.clear();

    return error;
}

void
XdnsHandler::Release(XdnsHandler* handler)
{
    if (handler != NULL)
    {
        delete handler;
    }
}

XdnsHandlerImpl::XdnsHandlerImpl(XdnsServerManager* serverManager)
: XdnsHandler(),
  m_serverManager(serverManager)
{
    assert(m_serverManager != NULL);
}

XdnsHandlerImpl::~XdnsHandlerImpl()
{
    if (m_serverManager != NULL)
    {
        XdnsServerManager::Release(m_serverManager);
    }
}

XdnsError
XdnsHandlerImpl::GetHostByName(const string& name,
                               XdnsListener* listener,
                               long* queryId)
{
    XdnsNameQuery* query = new (nothrow) XdnsNameQuery(name, listener);
    if (NULL == query)
    {
        return XE_NoMem;
    }

    assert(m_serverManager != NULL);
    XdnsError error = m_serverManager->SendQuery(query, queryId);
    if (error != XE_Ok)
    {
        delete query;
        return error;
    }

    return XE_Ok;
}

XdnsError
XdnsHandlerImpl::GetHostByAddress(const struct in_addr* ipAddress,
                                  XdnsListener* listener,
                                  long* queryId)
{
    if (NULL == ipAddress)
    {
        return XE_Inval;
    }

    XdnsAddressQuery* query = new (nothrow) XdnsAddressQuery(ipAddress, listener);
    if (NULL == query)
    {
        return XE_NoMem;
    }

    assert(m_serverManager != NULL);
    XdnsError error = m_serverManager->SendQuery(query, queryId);
    if (error != XE_Ok)
    {
        delete query;
        return error;
    }

    return XE_Ok;
}

void
XdnsHandlerImpl::CancelQuery(long queryId)
{
    assert(m_serverManager != NULL);
    m_serverManager->CancelQuery(queryId);
}

XdnsError
XdnsHandlerImpl::GetProcessingFds(int* nfds,
                                  fd_set* readfds,
                                  fd_set* writefds)
{
    assert(m_serverManager != NULL);
    return (m_serverManager->GetProcessingFds(nfds, readfds, writefds));
}

XdnsError
XdnsHandlerImpl::Process(fd_set* readfds,
                         fd_set* writefds)
{
    assert(m_serverManager != NULL);
    return (m_serverManager->Process(readfds, writefds));
}

void
XdnsHandlerImpl::ReplyReceived(long queryId,
                               XdnsError error,
                               const XdnsData* data,
                               size_t startOffset,
                               const XdnsHeaderInfo* header)
{
    // gupeng
    map<long, XdnsQuery*>::iterator queryIter = m_queries.find(queryId);
    if (queryIter == m_queries.end())
    {
        // The query is canceled.
        return;
    }
    XdnsQuery* query = queryIter->second;
    assert(query != NULL);

    // Retrieve the requested hostname or IP address.
    string requestedHostname;
    const struct in_addr* requestedIpAddress = NULL;
    XdnsQueryType queryType = query->GetType();
    switch (queryType)
    {
    case XQT_Hostname:
    {
        XdnsNameQuery* nameQuery = static_cast<XdnsNameQuery*>(query);
        requestedHostname = nameQuery->GetHostname();
        assert(!(requestedHostname.empty()));
        break;
    }
    case XQT_IpAddress:
    {
        XdnsAddressQuery* addressQuery = static_cast<XdnsAddressQuery*>(query);
        requestedIpAddress = addressQuery->GetAddress();
        assert(requestedIpAddress != NULL);
        break;
    }
    default:
    {
        assert(false);
        break;
    }
    }

    vector<XdnsHostEntry*> hostEntries;
    do
    {
        if (error != XE_Ok)
        {
            break;
        }

        // Check if the data is correct.
        if (NULL == data)
        {
            error = XE_Unexpected;
            break;
        }

        // Decode the header info if it is NULL.
        size_t offset = startOffset;
        XdnsHeaderInfo tempHeader;
        if (NULL == header)
        {
            error = XdnsDecoder::DecodeHeader(*data, offset, tempHeader);
            if (error != XE_Ok)
            {
                break;
            }
            header = &tempHeader;
        }
        assert(header != NULL);

        int answerCount = static_cast<int>(header->m_ancount);
        if ((header->m_qdcount != 1) || (answerCount <= 0))
        {
            error = XE_BadResponse;
            break;
        }

        // Decode the question info.
        offset += XDNS_HEADER_SIZE;
        XdnsQuestionInfo question;
        size_t questionSize = 0;
        error = XdnsDecoder::DecodeQuestion(*data,
                                            startOffset,
                                            offset,
                                            question,
                                            questionSize);
        if (error != XE_Ok)
        {
            break;
        }

        /* Check if the question info of response is same with the question
           info retrieved from the query.
        */
        XdnsQuestionInfo requestedQuestion;
        query->GetQuestion(requestedQuestion);
        if ((question.m_qtype != requestedQuestion.m_qtype) ||
            (question.m_qclass != requestedQuestion.m_qclass) ||
            (question.m_qname.compare(requestedQuestion.m_qname) != 0))
        {
            error = XE_BadResponse;
            break;
        }

        // Decode the answers of the resource record info.
        offset += questionSize;
        hostEntries.reserve(static_cast<size_t>(answerCount));
        while ((--answerCount) >= 0)
        {
            XdnsResourceRecordInfo resourceRecord;
            size_t resourceRecordSize = 0;
            error = XdnsDecoder::DecodeResourceRecord(*data,
                                                      startOffset,
                                                      offset,
                                                      resourceRecord,
                                                      resourceRecordSize);
            if (error != XE_Ok)
            {
                break;
            }

            // Generate the host entry.
            XdnsHostEntry* hostEntry = NULL;
            error = XdnsHostEntryGenerator::Generate(requestedIpAddress,
                                                     requestedHostname,
                                                     question,
                                                     resourceRecord,
                                                     *data,
                                                     startOffset,
                                                     &hostEntry);
            if (error != XE_Ok)
            {
                /* TODO: Some formats of DNS resource record are not
                         related with the host entry and not supported.
                */
                error = XE_Ok;
                continue;
            }

            assert(hostEntry != NULL);
            hostEntries.push_back(hostEntry);

            offset += resourceRecordSize;
        }
    }
    while (0);

    if ((XE_Ok == error) && (hostEntries.empty()))
    {
        error = XE_NoEnt;
    }

    query->NotifyListener(queryId,
                          error,
                          (XE_Ok == error) ? (&hostEntries) : NULL);

    // Destroy the host entries.
    vector<XdnsHostEntry*>::iterator hostEntryIter = hostEntries.begin();
    vector<XdnsHostEntry*>::iterator hostEntryEnd = hostEntries.end();
    for (; hostEntryIter != hostEntryEnd; ++hostEntryIter)
    {
        delete (*hostEntryIter);
    }
}

XdnsError
XdnsHandlerImpl::InitServerStates(const vector<XdnsServerConfig>& customServers,
                                  vector<XdnsServerState*>& serverStates)
{
    assert(serverStates.empty());

    XdnsInit* init = NULL;
    XdnsError error = XdnsInit::Init(&init);
    if (error != XE_Ok)
    {
        return error;
    }
    assert(init != NULL);

    // Retrieve the system name servers.
    vector<XdnsServerConfig> systemServers;
    init->GetSystemResolvers(systemServers);

    XdnsInit::Release(init);
    init = NULL;

    /* Append the system servers to the custom servers. The custom servers
       should have the higher priority than the system servers.
    */
    vector<XdnsServerConfig> allServers(customServers);
    allServers.insert(allServers.end(), systemServers.begin(), systemServers.end());

    vector<XdnsServerState*> newServerStates;
    vector<XdnsServerConfig> initializedServers;

    // Remove the duplicate servers and initialize the server states.
    vector<XdnsServerConfig>::iterator serverIter = allServers.begin();
    vector<XdnsServerConfig>::iterator serverEnd = allServers.end();
    for (; serverIter != serverEnd; ++serverIter)
    {
        const XdnsServerConfig& server = *serverIter;
        unsigned short port = server.m_serverPort;
        const string& addressString = server.m_serverIpAddress;
        vector<XdnsServerConfig>::iterator initializedIter = initializedServers.begin();
        vector<XdnsServerConfig>::iterator initializedEnd = initializedServers.end();
        for (; initializedIter != initializedEnd; ++initializedIter)
        {
            if ((0 == addressString.compare(initializedIter->m_serverIpAddress)) &&
                (initializedIter->m_serverPort == port))
            {
                break;
            }
        }

        // Check if the server is duplicate.
        if (initializedIter == initializedEnd)
        {
            // Convert the Internet host address from a string.
            struct in_addr serverAddress = {0};
            serverAddress.s_addr = inet_addr(addressString.c_str());
            if (INADDR_NONE == serverAddress.s_addr)
            {
                error = XE_Inval;
                break;
            }

            // Initialize a new server state.
            XdnsServerState* newState = NULL;
            error = XdnsServerState::Init(port,
                                          &serverAddress,
                                          &newState);
            if (error != XE_Ok)
            {
                break;
            }
            assert(newState != NULL);
            newServerStates.push_back(newState);

            // Save the server for further comparison.
            initializedServers.push_back(server);
        }
    }

    if (XE_Ok == error)
    {
        serverStates.swap(newServerStates);
        return XE_Ok;
    }

    // Destroy the server states for the returned error.
    vector<XdnsServerState*>::iterator stateIter = newServerStates.begin();
    vector<XdnsServerState*>::iterator stateEnd = newServerStates.end();
    for (; stateIter != stateEnd; ++stateIter)
    {
        XdnsServerState::Release(*stateIter);
    }
    newServerStates.clear();

    return error;
}

